fix(loss): count num_valid_samples from the sample mask, not the batch dim - #3850

…h dim

Six losses in this file derive num_valid_samples from the sample mask.
Three derive it from the batch dimension instead: MseValueLossFn from
values.shape[0], and both distillation losses from input_ids.shape[0].

That number is a gate, not just a log line. dtensor_policy_worker,
dtensor_policy_worker_v2 and dtensor_value_worker_v2 all do

    if num_valid_samples > 0:
        mb_losses.append(loss.item())
        all_mb_metrics.append(loss_metrics)

so a microbatch whose sample_mask is entirely zero -- which is what
overlong filtering, env-flagged masking and the critic's
seq_logprob_error_threshold filter produce -- reports the raw batch size,
passes the gate, and contributes its zero loss to the step's reported mean.
dpo.py and rm.py additionally use the value as a weighted-average
denominator.

DistillationLossFn keeps a batch-dimension fallback: its masking branch is
conditional, and with no mask every sample is valid by definition.

Tests are CPU-only -- they build the loss inputs directly rather than going
through prepare_loss_input, so no GPU and no distributed init. The
cross-tokenizer loss gets the same one-line change but no new test: every
CPU test for it calls the private helpers, and none reaches __call__ where
the metrics dict is built.

Signed-off-by: Tianyi Zhang <zhangtianyi975@gmail.com>
Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com>